package org.example.embeding;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;

/**
 * docker命令： docker run -d -p 6379:6379 --name redis-vector-db  redis/redis-stack-server
 *
 * @author fish_temp_author
 * @since fish_temp_since
 */
public class EmbeddingStoreWithRedis {
    public static void main(String[] args) {
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey("demo")
                .modelName("text-embedding-3-small")
                .build();

        TextSegment textSegment = TextSegment.textSegment("Java是世界上最好的语言！");
        Response<Embedding> embed = embeddingModel.embed(textSegment);
        System.out.println(embed.content().vector().length);
        System.out.println(embed.content());

        RedisEmbeddingStore redisEmbeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                // 通过特定的命令来使用，在redis-cli中：
                //                      FT._LIST
                //                      FT.INFO test222
                .indexName("test222")
                // key前缀
                .prefix("my:em:vec:")
                // 维度
                .dimension(embed.content().vector().length)
                .build();

        Embedding content = embed.content();
        // 指定id
        // redisEmbeddingStore.add(UUID.randomUUID().toString(), content);
        // 把原始内容也写进去
        redisEmbeddingStore.add(content, textSegment);

        TextSegment textSegment2 = TextSegment.textSegment("世界上最美的花是西兰花");
        redisEmbeddingStore.add(embeddingModel.embed(textSegment2).content(), textSegment2);

        TextSegment textSegment3 = TextSegment.textSegment("世界上最美的城市是杭州");
        redisEmbeddingStore.add(embeddingModel.embed(textSegment3).content(), textSegment2);

        EmbeddingSearchResult<TextSegment> searchResult = redisEmbeddingStore.search(EmbeddingSearchRequest.builder()
                .maxResults(3)
                .minScore(Double.valueOf("0.6"))
                .queryEmbedding(embeddingModel.embed("什么是世界上最好的语言？").content())
                .build());

        searchResult.matches().forEach(match -> {
            Double score = match.score();
            System.out.println(score);
            System.out.println(match.embedded().text());
        });
    }
}
